#include "main/utils.h"
using namespace std;

class Solution {
public:
    void sortColors(vector<int>& nums) {
      int size = nums.size();
      int idx0 = 0, idx2 = size - 1;
      for (int i = 0; i <= idx2; ++i) {
        if (nums[i] == 0) {
          swap(nums[idx0], nums[i]);
          ++idx0;
        } else if (nums[i] == 2) {
          swap(nums[i], nums[idx2]);
          --idx2;
          --i;
        }
      } 
    }

    void sortColors1(vector<int>& nums) {
      int n = nums.size();
      int idx0 = 0, idx1 = 0;
      for (int i = 0; i < n; ++i) {
        if (nums[i] == 0) {
          swap(nums[idx0], nums[i]);
          if (nums[i] == 1) {
            swap(nums[idx1], nums[i]);
            ++idx1;
          }
          ++idx0;
        } else if (nums[i] == 1) {
          if (idx1 < idx0)
            idx1 = idx0;
          swap(nums[idx1], nums[i]);
          ++idx1;
        }
      }
    }

    void sortColors2(vector<int>& nums) {
      int n = nums.size();
      int idx = 0;
      for (int i = 0; i < n; ++i) {
        if (nums[i] == 0) {
          swap(nums[idx], nums[i]);
          ++idx;
        }
      }
      for (int i = idx; i < n; ++i) {
        if (nums[i] == 1) {
          swap(nums[idx], nums[i]);
          ++idx;
        }
      }
    }
};

int main() {
  vector<int> nums = {2, 0, 2, 1, 1, 0};
  Solution su;
  su.sortColors(nums);
  printVec("leetcode75", nums);

  nums = {2, 0, 2, 1, 1, 0};
  su.sortColors1(nums);
  printVec("leetcode75", nums);

  nums = {2, 0, 2, 1, 1, 0};
  su.sortColors2(nums);
  printVec("leetcode75", nums);

  return 0;
}
